# ZITS/inference.py
import torch
from torchvision import transforms
from PIL import Image
from zits_model import ZITS  # 根据ZITS的代码实际引用
import numpy as np

class ZITSInpainter:
    def __init__(self, device="cuda"):
        self.device = device
        self.model = ZITS.from_pretrained("checkpoints/zits.pth")
        self.model.eval().to(device)

    def inpaint(self, image_bgr: np.ndarray, text: str) -> np.ndarray:
        # 转为RGB PIL
        image = Image.fromarray(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB))
        # 根据ZITS需求预处理并推理
        # 这里只是示意代码
        output = self.model.infer(image=image, prompt=text)
        out_np = np.array(output)
        return cv2.cvtColor(out_np, cv2.COLOR_RGB2BGR)
